import torch
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader


class Module(nn.Module):
    def __init__(self):
        super(Module,self).__init__()
        self.models = nn.Sequential(nn.Conv2d(in_channels=3,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=2),
                                    nn.Conv2d(in_channels=32,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Conv2d(in_channels=32,out_channels=64,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Flatten(),
                                    nn.Linear(1024,64),
                                    nn.Linear(64,10))
        # self.conv1 = nn.Conv2d(in_channels=3,out_channels=32,kernel_size=(5,5),padding=2)
        # self.pool1 = nn.MaxPool2d(kernel_size=2)
        # self.conv2 = nn.Conv2d(in_channels=32,out_channels=32,kernel_size=(5,5),padding=2)
        # self.pool2 = nn.MaxPool2d(kernel_size=(2,2))
        # self.conv3 = nn.Conv2d(in_channels=32,out_channels=64,kernel_size=(5,5),padding=2)
        # self.pool3 = nn.MaxPool2d(kernel_size=(2,2))
        # self.flatten = nn.Flatten()
        # self.linear1 = nn.Linear(1024,64)
        # self.linear2 = nn.Linear(64,10)
    def forward(self,x):
        #可以使用sequenti方法组成相关模型,简化代码
        # x = self.conv1(x)
        # x = self.pool1(x)
        # x = self.conv2(x)
        # x = self.pool2(x)
        # x = self.conv3(x)
        # x = self.pool3(x)
        # x = self.flatten(x)
        # x = self.linear1(x)
        # x = self.linear2(x)
        x = self.models(x)
        return x


data_test = torchvision.datasets.CIFAR10("./datavision", train=False, transform=torchvision.transforms.ToTensor())
datas = DataLoader(data_test,batch_size=64,shuffle=True,drop_last=True)
loss = nn.CrossEntropyLoss()
MyModule = Module()

for data in datas:
    imgs,targets = data
    output = MyModule(imgs)
    result_loss = loss(output,targets)
    result_loss.backward()
    print("ok")

